{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import json\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "pre_datas_test=\"/home/admin/xssdetection/file/pre_datas_test.csv\"\n",
    "vec_dir=\"/home/admin/xssdetection/file/word2vec.pickle\"\n",
    "def data_generator(data_dir):\n",
    "    reader = tf.TextLineReader()\n",
    "    queue = tf.train.string_input_producer([data_dir])\n",
    "    _, value = reader.read(queue)\n",
    "    coord = tf.train.Coordinator()\n",
    "    sess = tf.Session()\n",
    "    threads = tf.train.start_queue_runners(sess=sess, coord=coord)\n",
    "    while True:\n",
    "        v = sess.run(value)\n",
    "        [data, label] = v.split(b\"|\")\n",
    "        data = np.array(json.loads(data.decode(\"utf-8\")))\n",
    "        label = np.array(json.loads(label.decode(\"utf-8\")))\n",
    "        yield (data, label)\n",
    "    coord.request_stop()\n",
    "    coord.join(threads)\n",
    "    sess.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator=data_generator(pre_datas_test)\n",
    "data,label=next(generator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(532,)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_data = []\n",
    "data_embed = []\n",
    "with open(vec_dir,\"rb\") as f :\n",
    "    word2vec=pickle.load(f)\n",
    "    dictionary=word2vec[\"dictionary\"]\n",
    "    reverse_dictionary=word2vec[\"reverse_dictionary\"]\n",
    "    embeddings=word2vec[\"embeddings\"]\n",
    "for d in data:\n",
    "    if d != -1:\n",
    "        data_embed.append(embeddings[reverse_dictionary[d]])\n",
    "    else:\n",
    "        data_embed.append([0.0] * len(embeddings[\"UNK\"]))\n",
    "batch_data.append(data_embed)\n",
    "test_data=np.array(batch_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 532, 128)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(test_data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras.models import load_model\n",
    "model=load_model('./file/xss_lstm.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.99782795, 0.00217206], dtype=float32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#batch data(batch_size=500) testing,set steps=1\n",
    "model.predict(test_data,steps=1)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "#single data testing\n",
    "from keras.models import load_model\n",
    "model2=load_model('./file/xss_lstm_5000.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.9816626 , 0.01833741]], dtype=float32)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model2.predict(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
